// implement fork from user space
//#define _K_DEBUG_

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
static void
pgfault(struct UTrapframe *utf)
{
	void *addr = (void *) utf->utf_fault_va;
	uint32_t err = utf->utf_err;
	int r;


	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at vpt
	//   (see <inc/memlayout.h>).


	// LAB 4: Your code here.
        dbg_print("evid=0x%08x addr=0x%08x, err=0x%08x", sys_getenvid(), addr, err);
        if (!(err & FEC_WR)) {
                panic("pgfault: FEC_WR check error ");
        }

        if (!(vpt[VPN(addr)] & PTE_COW)) {
                panic("pgfault: PTE_COW check error ");
        }

	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.
	//   No need to explicitly delete the old page's mapping.
	
	// LAB 4: Your code here.
        if ((r=sys_page_alloc(0, PFTEMP, PTE_U | PTE_W | PTE_P)) < 0) {
                panic("pgfault: sys_page_alloc error %e ", r);
        }
        memmove(PFTEMP, ROUNDDOWN(addr, PGSIZE), PGSIZE);
        if ((r=sys_page_map(0, PFTEMP, 0, ROUNDDOWN(addr, PGSIZE), PTE_P | PTE_W | PTE_U))) {
                panic("pgfault: sys_page_map error %e ", r);
        }
        if ((r=sys_page_unmap(0, PFTEMP)) < 0) {
                panic("pgfault: sys_page_unmap error %e ", r);
        }
	
        //panic("pgfault not implemented");
}

//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why mark ours copy-on-write again
// if it was already copy-on-write?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
// 
int
duppage(envid_t envid, unsigned pn)
{
	int r;
	void *addr;
	pte_t pte;
                                        
        pte = vpt[pn];
        addr = (void*) (pn << PGSHIFT);
        dbg_print("addr=0x%08x, pn=0x%08x", addr, pn);
        if ((pte & PTE_W) || (pte & PTE_COW)) {
                if ((r=sys_page_map(0, addr, envid, addr, PTE_P | PTE_U | PTE_COW)) < 0) {
                        panic("sys_page_map fail: %e", r);
                }
                if ((r=sys_page_map(0, addr, 0, addr, PTE_P | PTE_U | PTE_COW)) < 0) {
                        panic("sys_page_map fail: %e", r);
                }
        } else {
                if ((r=sys_page_map(0, addr, envid, addr, PTE_P | PTE_U)) < 0) {
                        panic("sys_page_map fail: %e", r);
                }
        }

	// LAB 4: Your code here.
	//panic("duppage not implemented");
	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use vpd, vpt, and duppage.
//   Remember to fix "env" and the user exception stack in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
	// LAB 4: Your code here.
        int                   r, i, j, pn;  
	envid_t               envid;
	uint8_t              *addr;
	extern unsigned char  end[];
        extern volatile pte_t vpt[];     // VA of "virtual page table"
        extern volatile pde_t vpd[];     // VA of current page directory
        extern void _pgfault_upcall(void);


        set_pgfault_handler(pgfault);
        envid = sys_exofork();
        if (envid < 0) {
                panic("sys_exofork: %e", envid);
        }
        if (envid == 0) {
		env = &envs[ENVX(sys_getenvid())];
		return 0;
        }

        for (i=PDX(UTEXT); i <= PDX(USTACKTOP-PGSIZE); i++) {
                if (!(vpd[i] & PTE_P)) {
                        continue;
                }
                for (j=0; j < NPDENTRIES; j++) {
                        pn = (i<<10) + j;
                        if ((pn < VPN(UXSTACKTOP - PGSIZE)) && (vpt[pn] & PTE_P)) {
                                duppage(envid, pn);
                        }
                }
        }

        if ((r=sys_page_alloc(envid, (void*)(UXSTACKTOP-PGSIZE), 
                                        PTE_U | PTE_P | PTE_W)) < 0) {
                panic("sys_page_alloc fail: %e ", r); 
        }

        if ((r=sys_page_map(envid, (void*)(UXSTACKTOP-PGSIZE), 0, 
                                        (void*)UTEMP, PTE_U | PTE_P | PTE_W)) < 0) {
                panic("sys_page_map fail: %e ", r); 
        }

        memmove((void*)UTEMP, (void*)UXSTACKTOP-PGSIZE, PGSIZE);
        sys_page_unmap(0, UTEMP);

        if ((r=sys_env_set_pgfault_upcall(envid, (void *)_pgfault_upcall)) < 0) {
                panic("sys_env_set_pgfault_upcall fail: %e ", r);
        }

	// Start the child environment running
	if ((r = sys_env_set_status(envid, ENV_RUNNABLE)) < 0) {
		panic("sys_env_set_status: %e", r);
        }

        dbg_print("0x%08x => 0x%08x\n", sys_getenvid(), envid);
        return envid;
	//panic("fork not implemented");
}

// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}
